import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, roc_curve, auc
import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

def plot_training_curves(experiment_dir, fold):
    """绘制训练和验证的准确率、损失曲线"""
    # 读取日志文件
    log_file = os.path.join(experiment_dir, "train_log.csv")
    data = pd.read_csv(log_file)
    
    # 确保plots文件夹存在
    plots_dir = os.path.join(experiment_dir, "plots")
    os.makedirs(plots_dir, exist_ok=True)
    
    # 获取所有epoch数据
    epochs = data[data['Fold'] == fold]['Epoch'].unique()
    
    # 1. 绘制准确率曲线
    plt.figure(figsize=(10, 6))
    plt.plot(epochs, data[(data['Fold'] == fold)]['Val Accuracy'], label='Validation Accuracy')
    plt.plot(epochs, data[(data['Fold'] == fold)]['Train Accuracy'], label='Train Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title(f'Accuracy Curve (Fold {fold})')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(plots_dir, f'acc_curve_fold{fold}.png'), bbox_inches='tight')
    plt.close()
    
    # 2. 绘制损失曲线
    plt.figure(figsize=(10, 6))
    plt.plot(epochs, data[(data['Fold'] == fold)]['Val Loss'], label='Validation Loss')
    plt.plot(epochs, data[(data['Fold'] == fold)]['Train Loss'], label='Train Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title(f'Loss Curve (Fold {fold})')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(plots_dir, f'loss_curve_fold{fold}.png'), bbox_inches='tight')
    plt.close()

def plot_confusion_matrix(model, val_loader, device, experiment_dir, fold, num_classes=3):
    """使用最佳模型在验证集上绘制混淆矩阵"""
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for b, h, labels in val_loader:
            b, h, labels = b.to(device), h.to(device), labels.to(device)
            _, logits = model(b, h)
            preds = torch.argmax(logits, dim=1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    cm = confusion_matrix(all_labels, all_preds)
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=[f'Class {i}' for i in range(num_classes)],
                yticklabels=[f'Class {i}' for i in range(num_classes)])
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title(f'Confusion Matrix (Fold {fold})')
    
    plots_dir = os.path.join(experiment_dir, "plots")
    os.makedirs(plots_dir, exist_ok=True)
    plt.savefig(os.path.join(plots_dir, f'confusion_matrix_fold{fold}.png'), bbox_inches='tight')
    plt.close()

def plot_roc_curve(model, val_loader, device, experiment_dir, fold, num_classes=3):
    """使用最佳模型在验证集上绘制ROC曲线"""
    model.eval()
    all_probs = []
    all_labels = []
    
    with torch.no_grad():
        for b, h, labels in val_loader:
            b, h, labels = b.to(device), h.to(device), labels.to(device)
            _, logits = model(b, h)
            probs = torch.softmax(logits, dim=1)
            all_probs.extend(probs.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    all_probs = np.array(all_probs)
    all_labels = np.array(all_labels)
    
    # 计算每个类别的ROC曲线
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    
    for i in range(num_classes):
        fpr[i], tpr[i], _ = roc_curve((all_labels == i).astype(int), all_probs[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])
    
    # 计算微平均ROC曲线和AUC
    fpr["micro"], tpr["micro"], _ = roc_curve(all_labels.ravel(), all_probs.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])
    
    # 绘制所有ROC曲线
    plt.figure(figsize=(8, 8))
    colors = ['aqua', 'darkorange', 'cornflowerblue']
    for i, color in zip(range(num_classes), colors):
        plt.plot(fpr[i], tpr[i], color=color, lw=2,
                 label=f'ROC curve of class {i} (area = {roc_auc[i]:0.2f})')
    
    plt.plot([0, 1], [0, 1], 'k--', lw=2)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(f'ROC Curve (Fold {fold})')
    plt.legend(loc="lower right")
    
    plots_dir = os.path.join(experiment_dir, "plots")
    os.makedirs(plots_dir, exist_ok=True)
    plt.savefig(os.path.join(plots_dir, f'roc_curve_fold{fold}.png'), bbox_inches='tight')
    plt.close()